import cv2
import numpy as np
from .preprocessing import Preprocessing


class ImagePreprocessor(Preprocessing):
    def __init__(self, target_size=(640, 640)):
        self.target_size = target_size

    def preprocess(self, image_bytes: bytes) -> np.ndarray:
        """
        预处理图像（以JPEG或PNG编码），为Triton推断做准备。

        参数:
            image_bytes: 编码后的图像数据。

        返回值:
            np.ndarray: 经过预处理的图像，以numpy数组形式返回，数组形状为(height, width, channels)。
        """
        # 使用OpenCV从缓冲区解码图像
        image = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_COLOR)

        # 执行必要的图像调整大小、标准化等操作
        resized_image = self._resize_and_normalize(image)

        return resized_image

    def _resize_and_normalize(self, image: np.ndarray) -> np.ndarray:
        """
        将图像调整到目标大小，并转换为BGR格式。

        参数:
            image: 输入的图像，以numpy数组形式表示。

        返回值:
            np.ndarray: 调整大小并转换为BGR格式的图像。
        """
        resized_image = cv2.resize(image, self.target_size)
        return resized_image


# 使用示例
# preprocessor = ImagePreprocessor()
# image_bytes = ...  # 获取图像字节数据
# preprocessed_image = preprocessor.preprocess(image_bytes)
